# -*- coding: utf-8 -*-
"""
Created on Nov 18 22:57:41 2022

@author: zhaoxm
"""
import numpy as np
import cv2
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm

data_path = "D:/data/diamond_dataset/images/"
save_path = "D:/data/diamond_dataset/squared_images/"
save_format = "square"
#%%
save_path = Path(save_path)
save_path.mkdir(exist_ok=True)
ls = sorted(Path(data_path).glob("*.png"))

widths, heights = [], []
for img_path in tqdm(ls):
    image = cv2.imread(str(img_path))
    # image = cv2.pyrMeanShiftFiltering(image, 10, 100)
    # x, thr = cv2.threshold(gray, 170, 255, cv2.THRESH_BINARY)
    # plt.imshow(gray)
    
    height, width = image.shape[:2]
    
    # widths.append(width)
    # heights.append(height)
    
    if save_format == "resize":
    
        edge = int((640 * height / 576 - width) / 2)
        xmap, ymap = np.meshgrid(np.linspace(-edge, width - 1 + edge, 640, dtype=np.float32), np.linspace(0, height - 1, 576, dtype=np.float32))
        resize_image = cv2.remap(image, xmap, ymap, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
        
        # fig, ax = plt.subplots(1, 2)
        # ax[0].imshow(image)
        # ax[1].imshow(resize_image)
        # plt.show()
        
        cv2.imwrite(str(save_path / img_path.name), resize_image)
    
    elif save_format == "square":
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        output = image.copy()
        center = np.array([width/2, height/2])
        circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1.0, 500, param1=50, param2=30, minRadius=20, maxRadius=0)
        assert circles is not None
        if circles is not None:
            circles = np.round(circles[0, :]).astype("int")
            idx = np.argmin(np.linalg.norm(circles[:, :2] - center, axis=1))
            x, y, r = circles[idx]
            # for (x, y, r) in circles:
            cv2.circle(output, (x, y), r, (0, 255, 0), 4)
            cv2.rectangle(output, (x - 5, y - 5), (x + 5, y + 5), (255, 128, 0), -1)
            # cv2.imshow("output", np.hstack([image, output]))
            # cv2.waitKey(0)  
            # plt.imshow(np.hstack([image, output]))
            # plt.show()
            
            xmap, ymap = np.meshgrid(np.linspace(x-r, x+r, 640, dtype=np.float32),
                                     np.linspace(y-r, y+r, 640, dtype=np.float32))
            resize_image = cv2.remap(image, xmap, ymap, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
            # plt.imshow(resize_image)
            # plt.show()
            
            cv2.imwrite(str(save_path / img_path.name), resize_image)
        
        # cv2.destroyAllWindows()
